#include "CSVExporter.h"

#include "LocalTrajExtraction.h"
#include <MMM/Motion/Plugin/KinematicPlugin/KinematicSensor.h>
#include <MMM/Motion/Plugin/KinematicPlugin/KinematicSensorMeasurement.h>
#include <MMM/Motion/Plugin/ModelPosePlugin/ModelPoseSensor.h>
#include <SimoxUtility/algorithm/string.h>
#include <SimoxUtility/math/convert.h>
#include <MMM/Motion/MotionReaderXML.h>
#include <VirtualRobot/Robot.h>
#include <VirtualRobot/RobotNodeSet.h>
#include <MMMSimoxTools/MMMSimoxTools.h>
#include <MMM/Model/Model.h>

#include <MMM/Exceptions.h>

using namespace MMM;

void CSVExportPlugin::exportCSV_dir(const std::string &directoryPath, const std::string &exportDirectoryPath, bool localFrame, bool exportHands) {
    for(auto &file : std::filesystem::directory_iterator(directoryPath)) {
        auto fileExtension = file.path().extension();
        if (fileExtension == ".xml") {
            exportCSV(file.path(), exportDirectoryPath, localFrame, exportHands);
        }
    }
}

void CSVExportPlugin::exportCSV(const std::string &motionPath, const std::string &exportDirectoryPath, bool localFrame, bool exportHands) {
    try {
        MMM_INFO << "Reading motion from " << motionPath << std::endl;
        MMM::MotionList motions = motionReader->loadAllMotions(motionPath);
        auto timesteps = MMM::Motion::calculateMinMaxTimesteps(motions);
        exportCSV(motions, exportDirectoryPath, localFrame, exportHands, std::get<0>(timesteps), std::get<1>(timesteps), 100);

    } catch (MMM::Exception::MMMException& e) {
        MMM_ERROR << e.what() << endl;
    }
}

void CSVExportPlugin::exportCSV(MotionList motions, const std::string &exportDirectoryPath, bool localFrame, bool exportHands, float minTimestep, float maxTimestep, float framesPerSecond) {
    if (localFrame){
        std::vector<std::string> objNames;
        for(auto & motion : motions)
        {
            if (!motion->getModel()) continue;
            std::string s = std::filesystem::path(motion->getModel()->getFilename()).stem();
            objNames.push_back(s);
        }
        objNames.push_back("MMMHandL");
        objNames.push_back("MMMHandR");

        MMM_INFO << "Calculating local trajectories";
        MMM::MotionList motions_orig = motions;
        for(unsigned int j = 0; j < objNames.size(); j++){
            for(unsigned int k = 0; k < objNames.size(); k++){
                if(j != k){
                    MotionPtr new_motion = nullptr;
                    LocalTrajExtractionPtr locTrajEx(new LocalTrajExtraction(motions_orig, minTimestep, maxTimestep, framesPerSecond, objNames[j], objNames[k]));
                    new_motion = locTrajEx->getLocalPoseTraj();
                    motions.push_back(new_motion);
                }
            }
        }
    }

    unsigned int size = motions.size();
    for(unsigned int i = 0; i < size; i++)
    {
        auto motion = motions.at(i);

        //if (!motion->getModel()) continue;
        if (!motion->hasSensor("ModelPose") && !motion->hasSensor("Kinematic")) {
            MMM_INFO << "Not exporting motion with name '" << motion->getName() << "'. The motion has no model pose or kinematic sensor!" << std::endl;
            continue;
        }
        std::vector<std::string> headers{"timestep"};
        std::vector<std::vector<float>> csvData;
        auto poseSensor = motion->getSensorByType<ModelPoseSensor>("ModelPose");
        if (poseSensor)
        {
            auto names = {"RootPosition_X",
                         "RootPosition_Y",
                         "RootPosition_Z",
                         "RootOrientation_qx",
                         "RootOrientation_qy",
                         "RootOrientation_qz",
                         "RootOrientation_qw"};
            headers.insert(headers.end(), names.begin(), names.end());
        }
        auto kinSensors = motion->getSensorsByType<KinematicSensor>("Kinematic");
        for (auto kinSensor : kinSensors)
        {
            auto names = kinSensor->getJointNames();
            headers.insert(headers.end(), names.begin(), names.end());
        }

        for (float t = minTimestep; t <= maxTimestep; t += 1.0/framesPerSecond)
        {
            std::vector<float> timestepData;
            timestepData.push_back(t);

            if(poseSensor)
            {
                auto data = boost::dynamic_pointer_cast<ModelPoseSensorMeasurement>(poseSensor->getMeasurement(t));
                Eigen::Vector3f pos = data->getRootPosition();
                timestepData.push_back(pos(0));
                timestepData.push_back(pos(1));
                timestepData.push_back(pos(2));
                Eigen::Quaternionf rot = simox::math::rpy_to_quat(data->getRootRotation());
                timestepData.push_back(rot.x());
                timestepData.push_back(rot.y());
                timestepData.push_back(rot.z());
                timestepData.push_back(rot.w());
            }

            for (auto kinSensor : kinSensors)
            {
                auto data = boost::dynamic_pointer_cast<KinematicSensorMeasurement>(kinSensor->getMeasurement(t));
                Eigen::VectorXf angles = data->getJointAngles();
                timestepData.insert(timestepData.end(), angles.data(), angles.data()+ angles.rows());
            }

            csvData.push_back(timestepData);
        }

        auto folderPath = std::filesystem::path(exportDirectoryPath) / std::filesystem::path(motion->getOriginFilePath()).stem();
        if (!std::filesystem::exists(folderPath))
            std::filesystem::create_directory(folderPath);

        std::string filename = folderPath / std::filesystem::path(motion->getName() + ".csv");
        MMM_INFO << "Writing to " << filename << std::endl;
        std::ofstream f(filename.c_str());
        f << boost::algorithm::join(headers, ",") << "\n";
        for(auto & entry: csvData)
        {
            std::vector<std::string> valueStrings;
            for(auto v : entry)
            {
                valueStrings.push_back(simox::alg::to_string(v));
            }
            f <<  boost::algorithm::join(valueStrings, "," ) << "\n";
        }

        if (exportHands && std::filesystem::path(motion->getModel()->getFilename()).stem() == "mmm") {
            VirtualRobot::RobotPtr robot = MMM::SimoxTools::buildModel(motion->getModel(), false);
            std::vector<VirtualRobot::RobotNodeSetPtr> rns;
            for (auto kinSensor : kinSensors)
            {
                rns.push_back(VirtualRobot::RobotNodeSet::createRobotNodeSet(robot, "RobotNodeSet", kinSensor->getJointNames(), "", "", true));
            }

            // empty model should be no problem
            ModelPtr emptyModel(new Model());
            MotionPtr leftHandMotion(new Motion(motion->getName() + "_LeftHand", emptyModel, emptyModel, nullptr, motion->getOriginFilePath()));
            ModelPoseSensorPtr modelPoseSensor_left(new ModelPoseSensor());
            leftHandMotion->addSensor(modelPoseSensor_left);
            MotionPtr rightHandMotion(new Motion(motion->getName() + "_RightHand", emptyModel, emptyModel, nullptr, motion->getOriginFilePath()));
            ModelPoseSensorPtr modelPoseSensor_right(new ModelPoseSensor());
            rightHandMotion->addSensor(modelPoseSensor_right);

            for (float t = minTimestep; t <= maxTimestep; t += 1.0/framesPerSecond)
            {
                if (poseSensor)
                    robot->setGlobalPose(poseSensor->getDerivedMeasurement(t)->getRootPose());

                for (unsigned int i = 0; i < kinSensors.size(); i++)
                {
                    rns[i]->setJointValues(kinSensors[i]->getDerivedMeasurement(t)->getJointAngles());
                }

                                modelPoseSensor_left->addSensorMeasurement(ModelPoseSensorMeasurementPtr(new ModelPoseSensorMeasurement(t, robot->getRobotNode("LWsegment_joint")->getGlobalPose())));
                                modelPoseSensor_right->addSensorMeasurement(ModelPoseSensorMeasurementPtr(new ModelPoseSensorMeasurement(t, robot->getRobotNode("RWsegment_joint")->getGlobalPose())));
                //modelPoseSensor_left->addSensorMeasurement(ModelPoseSensorMeasurementPtr(new ModelPoseSensorMeasurement(t, robot->getRobotNode("LWsegment_joint")->getPoseInRootFrame())));
                //modelPoseSensor_right->addSensorMeasurement(ModelPoseSensorMeasurementPtr(new ModelPoseSensorMeasurement(t, robot->getRobotNode("RWsegment_joint")->getPoseInRootFrame())));
            }
            motions.push_back(leftHandMotion);
            motions.push_back(rightHandMotion);
            size += 2;
        }
    }
}

